from multiprocessing import Process
import numpy as np
import os
import scipy
from scipy.io import loadmat
from scipy.ndimage import gaussian_filter
import matplotlib.pyplot as plt
import glob

# 进程数
mulpNumber = 4
# 根目录
root = 'D:/STUDY/CrowdCounting/MCNN-pytorch'

part_A_train = os.path.join(root, 'part_A/train_data', 'images')
part_B_train = os.path.join(root, 'part_B/train_data', 'images')
part_A_test = os.path.join(root, 'part_A/test_data', 'images')
part_B_test = os.path.join(root, 'part_B/test_data', 'images')


def gaussian_filter_density(img, points):
    '''
    This code use k-nearst, will take one minute or more to generate a density-map with one thousand people.

    points: a two-dimension list of pedestrians' annotation with the order [[col,row],[col,row],...].
    img_shape: the shape of the image, same as the shape of required density-map. (row,col). Note that can not have channel.

    return:
    density: the density-map we want. Same shape as input image but only has one channel.

    example:
    points: three pedestrians with annotation:[[163,53],[175,64],[189,74]].
    img_shape: (768,1024) 768 is row and 1024 is column.
    '''
    # 获取图片的h和我
    img_shape = [img.shape[0], img.shape[1]]
    print("Shape of current image: ", img_shape, ". Totally need generate ", len(points), "gaussian kernels.")
    # 生成与图片相同大小的0矩阵图
    density = np.zeros(img_shape, dtype=np.float32)
    # 获取人数数据
    gt_count = len(points)
    # 如果人数数据为0直接密度图
    if gt_count == 0:
        return density

    leafsize = 2048
    # build kdtree

    tree = scipy.spatial.KDTree(points.copy(), leafsize=leafsize)
    # query kdtree
    distances, locations = tree.query(points, k=4)

    print('generate density...')
    for i, pt in enumerate(points):
        pt2d = np.zeros(img_shape, dtype=np.float32)
        if int(pt[1]) < img_shape[0] and int(pt[0]) < img_shape[1]:
            pt2d[int(pt[1]), int(pt[0])] = 1.
        else:
            continue
        if gt_count > 1:
            sigma = (distances[i][1] + distances[i][2] + distances[i][3]) * 0.1
        else:
            sigma = np.average(np.array(gt.shape)) / 2. / 2.  # case: 1 point
        density += scipy.ndimage.gaussian_filter(pt2d, sigma, mode='constant')
    print('done.')
    return density


def fun1(_img_paths_, number):
    print("第", number + 1, "号进程开启")
    img_paths_ = _img_paths_

    for img_path in img_paths_:
        matName = img_path.replace('jpg', 'mat').replace('images', 'ground_truth').replace('IMG_', 'GT_IMG_')
        mat = loadmat(img_path.replace('jpg', 'mat').replace('images', 'ground_truth').replace('IMG_', 'GT_IMG_'))
        img = plt.imread(img_path)  # 768行*1024列
        k = np.zeros((img.shape[0], img.shape[1]))
        points = mat["image_info"][0, 0][0, 0][0]  # 1546person*2(col,row)
        k = gaussian_filter_density(img, points)
        # plt.imshow(k,cmap=CM.jet)
        # save density_map to disk
        np.save(img_path.replace('.jpg', '.npy').replace('images', 'ground_truth'), k)
    return


def list_split(items, n):
    return [items[i:i + n] for i in range(0, len(items), n)]


if __name__ == '__main__':

    path_sets = [part_A_train, part_B_train, part_A_test, part_B_test]

    img_paths = []
    for path in path_sets:
        for img_path in glob.glob(os.path.join(path, '*.jpg')):
            img_paths.append(img_path)
    print(len(img_paths))
    splitNumber = int(len(img_paths) / mulpNumber) + 1
    paths = [img_paths[i:i + splitNumber] for i in range(0, len(img_paths), splitNumber)]

    process_list = []
    for i in range(mulpNumber):  # 开启5个子进程执行fun1函数
        p = Process(target=fun1, args=(paths[i], i))  # 实例化进程对象
        p.start()
        process_list.append(p)

    for i in process_list:
        p.join()

    print('计算结束')

"""
    process_list = []
    for i in range(4):
        p = Process(target=fun1, args=(i,))  # 实例化进程对象
        p.start()
        process_list.append(p)

    for i in process_list:
        p.join()

    print('计算结束')

"""
